package com.jair.jairaiagent.rag;

import jakarta.annotation.Resource;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.pgvector.PgVectorStore;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.jdbc.core.JdbcTemplate;

import javax.sql.DataSource;
import java.util.List;
import static org.springframework.ai.vectorstore.pgvector.PgVectorStore.PgDistanceType.COSINE_DISTANCE;
import static org.springframework.ai.vectorstore.pgvector.PgVectorStore.PgIndexType.HNSW;

@Configuration
public class PgVectorVectorStoreConfig {

    @Resource
    private LoveAppDocumentLoder loveAppDocumentLoader;


    @Bean
    public VectorStore pgVectorVectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel dashscopeEmbeddingModel) {
        VectorStore vectorStore = PgVectorStore.builder(jdbcTemplate, dashscopeEmbeddingModel)
                .dimensions(1536)                    // Optional: defaults to model dimensions or 1536
                .distanceType(COSINE_DISTANCE)       // Optional: defaults to COSINE_DISTANCE
                .indexType(HNSW)                     // Optional: defaults to HNSW
                .initializeSchema(true)              // Optional: defaults to false
                .schemaName("public")                // Optional: defaults to "public"
                .vectorTableName("vector_store")     // Optional: defaults to "vector_store"
                .maxDocumentBatchSize(10000)         // Optional: defaults to 10000
                .build();
         // 测试文档加载
         try {
            System.out.println("开始加载文档...");
            List<Document> documents = loveAppDocumentLoader.loadMarkDowns();
            System.out.println("文档加载完成，共加载 " + (documents != null ? documents.size() : 0) + " 个文档");
            
            if (documents != null && !documents.isEmpty()) {
                System.out.println("开始添加文档到向量存储...");
                vectorStore.add(documents);
                System.out.println("文档添加完成");
            } else {
                System.out.println("没有文档需要添加");
            }
        } catch (Exception e) {
            System.err.println("文档加载或添加过程中出现错误: " + e.getMessage());
            e.printStackTrace();
        }
        return vectorStore;
    }
}
